import os
import pty
import select
import sys
import termios
import tty

class PTY:
    def __init__(self):
        # 创建伪终端
        self.master_fd, self.slave_fd = pty.openpty()
        self.child_pid = os.fork()
        
        if self.child_pid == 0:
            # 子进程：设置伪终端为原始模式并执行 shell
            os.close(self.master_fd)
            slave_name = os.ttyname(self.slave_fd)
            os.setsid()
            os.dup2(self.slave_fd, 0)
            os.dup2(self.slave_fd, 1)
            os.dup2(self.slave_fd, 2)
            os.close(self.slave_fd)
            os.execv('/bin/bash', ['/bin/bash'])
            # "docker run -it -v $(pwd):/workdir registry.cn-wulanchabu.aliyuncs.com/sky-group/joern:20250418 joern"
            # os.execv('docker', ['docker', "run", "-it", "-v", os.getcwd()+":/workdir", "registry.cn-wulanchabu.aliyuncs.com/sky-group/joern:20250418", "joern"])
        else:
            # 父进程：设置伪终端为原始模式
            os.close(self.slave_fd)
            tty.setraw(self.master_fd)

    def read_output(self, timeout=0.1):
        """
        读取终端输出。
        :param timeout: 超时时间（秒）
        :return: 终端输出字符串
        """
        output = ""
        while True:
            readable, _, _ = select.select([self.master_fd], [], [], timeout)
            if readable:
                data = os.read(self.master_fd, 1024).decode('utf-8')
                if not data:
                    break
                output += data
            else:
                break
        return output

    def write_command(self, command):
        """
        在终端中注入命令。
        :param command: 要注入的命令字符串
        """
        os.write(self.master_fd, (command + "\n").encode('utf-8'))

    def close(self):
        """
        关闭伪终端。
        """
        os.close(self.master_fd)
        os.kill(self.child_pid, 15)  # 发送 SIGTERM 信号终止子进程

# 示例用法
if __name__ == "__main__":
    pty_instance = PTY()
    try:
        # while 
        # pty_instance.write_command()
        print(pty_instance.read_output())
        while True:
            pty_instance.write_command(input())
            print(pty_instance.read_output())
    finally:
        pty_instance.close()